//
// Created by Administrator on 2021/6/11.
//

#include "NvInferPlugin.h"
#include "dbnet.h"
#include "utils/clipper.hpp"
#include "utils/RrConfig.h"
#include "utils/utils.h"
//#include "common.h"
//#define CHECK(exp) \
//if (exp != cudaSuccess) { \
//printf("[ERROR] [%s:%d] %d\n", __FILE__, __LINE__, cudaGetLastError()); \
//exit(-1); \
//}
//
//class Logger : public nvinfer1::ILogger {
//public:
//
//    Logger() : Logger(Severity::kWARNING) {}
//
//    Logger(Severity severity) : reportableSeverity(severity) {}
//
//    void log(Severity severity, const char *msg) override {
//        // suppress messages with severity enum value greater than the reportable
//        if (severity > reportableSeverity) return;
//
//        switch (severity) {
//            case Severity::kINTERNAL_ERROR:
//                std::cerr << "INTERNAL_ERROR: ";
//                break;
//            case Severity::kERROR:
//                std::cerr << "ERROR: ";
//                break;
//            case Severity::kWARNING:
//                std::cerr << "WARNING: ";
//                break;
//            case Severity::kINFO:
//                std::cerr << "INFO: ";
//                break;
//            default:
//                std::cerr << "UNKNOWN: ";
//                break;
//        }
//        std::cerr << msg << std::endl;
//    }
//
//    Severity reportableSeverity{Severity::kWARNING};
//};

using IOFormat = std::pair<nvinfer1::DataType, nvinfer1::TensorFormats>;

inline bool broadcastIOFormats(const std::vector<IOFormat>& formats, size_t nbBindings, bool isInput = true)
{
    bool broadcast = formats.size() == 1;
    bool validFormatsCount = broadcast || (formats.size() == nbBindings);
    if (!formats.empty() && !validFormatsCount)
    {
        if (isInput)
        {
            throw std::invalid_argument(
                    "The number of inputIOFormats must match network's inputs or be one for broadcasting.");
        }
        else
        {
            throw std::invalid_argument(
                    "The number of outputIOFormats must match network's outputs or be one for broadcasting.");
        }
    }
    return broadcast;
}

int DBnet::serialization(const std::string &model_file, const std::string &engine_file) {

    Logger gLogger(nvinfer1::ILogger::Severity::kERROR);//kWARNING, kINFO, kINTERNAL_ERROR

    int device_count = 0;
    cudaGetDeviceCount(&device_count);
    assert(device_count > 0);
    if ((device_id + 1) > device_count) {
        device_id = 0;
    }
    std::cout << "GPU_nums: " << device_count << std::endl;
    std::cout << "Dbnet GPU ID: " << device_id << std::endl;
    cudaSetDevice(device_id);

    initLibNvInferPlugins(&gLogger, "");

//    std::unique_ptr <nvinfer1::IRuntime, nvinfer1::Destroy<nvinfer1::IRuntime>> runtime{
//            nvinfer1::createInferRuntime(gLogger)};

    std::unique_ptr <nvinfer1::IBuilder, nvinfer1::Destroy<nvinfer1::IBuilder>> builder{
            nvinfer1::createInferBuilder(gLogger)};
    if (!builder) {
        return -1;
    }

    const auto explicitBatch = 1U << static_cast<uint32_t>(nvinfer1::NetworkDefinitionCreationFlag::kEXPLICIT_BATCH);
    std::unique_ptr <nvinfer1::INetworkDefinition, nvinfer1::Destroy<nvinfer1::INetworkDefinition>> network{
            builder->createNetworkV2(explicitBatch)};
    if (!network) {
        return -1;
    }

    std::unique_ptr <nvonnxparser::IParser, nvinfer1::Destroy<nvonnxparser::IParser>> parser{
            nvonnxparser::createParser(*network, gLogger)};
    if (!parser) {
        return -1;
    }

    std::unique_ptr <nvinfer1::IBuilderConfig, nvinfer1::Destroy<nvinfer1::IBuilderConfig>> config{
            builder->createBuilderConfig()};
    if (!config) {
        return -1;
    }

    auto parsed = parser->parseFromFile(model_file.c_str(), static_cast<int>(nvinfer1::ILogger::Severity::kERROR));
    if (!parsed) {
        std::cout << "ERROR: could not parse input engine." << std::endl;
        return -1;
    }

    builder->setMaxWorkspaceSize(((1ULL << 10) * 1));

    {
        builder->setFp16Mode(true);
        builder->setMaxBatchSize(2);
        auto profile = builder->createOptimizationProfile();
        profile->setDimensions(network->getInput(0)->getName(), nvinfer1::OptProfileSelector::kMIN,
                               nvinfer1::Dims4{1, 3, 320, 320});
        profile->setDimensions(network->getInput(0)->getName(), nvinfer1::OptProfileSelector::kOPT,
                               nvinfer1::Dims4{1, 3, 1024, 1024});
        profile->setDimensions(network->getInput(0)->getName(), nvinfer1::OptProfileSelector::kMAX,
                               nvinfer1::Dims4{2, 3, 1920, 1920});
        config->addOptimizationProfile(profile);
        config->setFlag(nvinfer1::BuilderFlag::kFP16);
    }
    config->setMaxWorkspaceSize(static_cast<size_t>(1) << 10);
//    config->setMaxWorkspaceSize(static_cast<size_t>(16) << 20);
    config->setFlag(nvinfer1::BuilderFlag::kDISABLE_TIMING_CACHE);
    std::unique_ptr <nvinfer1::ICudaEngine, nvinfer1::Destroy<nvinfer1::ICudaEngine>> serialize_eigine{builder->buildEngineWithConfig(*network, *config)};

    // 序列化
    std::unique_ptr <nvinfer1::IHostMemory, nvinfer1::Destroy<nvinfer1::IHostMemory>> trtModelStream{
            serialize_eigine->serialize()};
    nvinfer1::writeBuffer(trtModelStream->data(), trtModelStream->size(), engine_file.c_str());

    return 0;
}

int DBnet::init(const std::string &model_file) {
    Logger gLogger(nvinfer1::ILogger::Severity::kERROR);//kWARNING, kINFO, kINTERNAL_ERROR
    int device_count = 0;
    cudaGetDeviceCount(&device_count);
    assert(device_count > 0);
    if ((device_id + 1) > device_count) {
        device_id = 0;
    }
    std::cout << "GPU_nums: " << device_count << std::endl;
    std::cout << "Dbnet GPU ID: " << device_id << std::endl;
    cudaSetDevice(device_id);

    //反序列化
    size_t dot_pos = model_file.find(".");
    std::string engine_file = model_file.substr(0, dot_pos) + ".engine";
    std::string buffer = nvinfer1::readBuffer(engine_file.c_str());

    if (trt_serialize) {
        if (!buffer.size()){
            std::string TrtRoot = RrConfigMgr::GetInstance().ReadString("ocrConfig", "TrtRoot", "");
            std::string name = TrtRoot + "/bin/trtexec --onnx=";
            //name += model_file + " --tacticSources=-cublasLt,+cublas --workspace=1024 ";
            name += model_file + " --workspace=1024 ";
            name += "--minShapes=input_0:1x3x320x320 --optShapes=input_0:1x3x1120x1120 --maxShapes=input_0:1x3x1920x1920 ";
            name += "--fp16 --saveEngine=" + engine_file;
            system(name.c_str());
        }
    } else {
        if (!buffer.size()) {
            serialization(model_file, engine_file);
        }
    }
    std::unique_ptr <nvinfer1::IRuntime, nvinfer1::Destroy<nvinfer1::IRuntime>> mRuntime{
            nvinfer1::createInferRuntime(gLogger)};
    buffer = nvinfer1::readBuffer(engine_file.c_str());
    mEngine.reset(mRuntime->deserializeCudaEngine(buffer.data(), buffer.size(), nullptr));
    context.reset(mEngine->createExecutionContext());
    assert(mEngine->getNbBindings() == MAX_BINDING_NUMS);
    cudaStreamCreate(&stream_t);

    return 0;
}

float paddimg(cv::Mat &In_Out_img, int shortsize = 960) {
    int w = In_Out_img.cols;
    int h = In_Out_img.rows;
    float scale = 1.f;
    if (w < h) {
        scale = (float) shortsize / w;
        h = scale * h;
        w = shortsize;
    } else {
        scale = (float) shortsize / h;
        w = scale * w;
        h = shortsize;
    }

    if (h % 32 != 0) {
        h = (h / 32 + 1) * 32;
    }
    if (w % 32 != 0) {
        w = (w / 32 + 1) * 32;
    }

//    cv::resize(In_Out_img, In_Out_img, cv::Size(w, h));
    cv::resize(In_Out_img, In_Out_img, cv::Size(shortsize, shortsize));
    return scale;
}

bool DBnet::Input(const cv::Mat &src, float *data) {
    int i = 0;
//    int batch = 0;
    std::vector<float> mean_value{0.406, 0.456, 0.485};  // BGR
    std::vector<float> std_value{0.225, 0.224, 0.229};
    for (int row = 0; row < src.rows; ++row) {
        uchar *uc_pixel = src.data + row * src.step;
        for (int col = 0; col < src.cols; ++col) {
            data[i] = (uc_pixel[2] / 255.0 - mean_value[2]) / std_value[2];
            data[i + src.rows * src.cols] = (uc_pixel[1] / 255.0 - mean_value[1]) / std_value[1];
            data[i + 2 * src.rows * src.cols] = (uc_pixel[0] / 255.0 - mean_value[0]) / std_value[0];
            uc_pixel += 3;
            ++i;
        }
    }
    return true;
}

float get_box_score(std::vector<float> &map, cv::Point2f rect[], int width, int height,
                    float threshold) {

    int xmin = width - 1;
    int ymin = height - 1;
    int xmax = 0;
    int ymax = 0;

    for (int j = 0; j < 4; j++) {
        if (rect[j].x < xmin) {
            xmin = rect[j].x;
        }
        if (rect[j].y < ymin) {
            ymin = rect[j].y;
        }
        if (rect[j].x > xmax) {
            xmax = rect[j].x;
        }
        if (rect[j].y > ymax) {
            ymax = rect[j].y;
        }
    }
    float sum = 0;
    int num = 0;
    for (int i = ymin; i <= ymax; i++) {
        for (int j = xmin; j <= xmax; j++) {
            if (map[i * width + j] > threshold) {
                sum = sum + map[i * width + j];
                num++;
            }
        }
    }

    return sum / num;
}

//bool get_mini_boxes(cv::RotatedRect &rotated_rect, cv::Point2f rect[],
//                    int min_size) {
//
//    cv::Point2f temp_rect[4];
//    rotated_rect.points(temp_rect);
//    for (int i = 0; i < 4; i++) {
//        for (int j = i + 1; j < 4; j++) {
//            if (temp_rect[i].x > temp_rect[j].x) {
//                cv::Point2f temp;
//                temp = temp_rect[i];
//                temp_rect[i] = temp_rect[j];
//                temp_rect[j] = temp;
//            }
//        }
//    }
//    int index0 = 0;
//    int index1 = 1;
//    int index2 = 2;
//    int index3 = 3;
//    if (temp_rect[1].y > temp_rect[0].y) {
//        index0 = 0;
//        index3 = 1;
//    } else {
//        index0 = 1;
//        index3 = 0;
//    }
//    if (temp_rect[3].y > temp_rect[2].y) {
//        index1 = 2;
//        index2 = 3;
//    } else {
//        index1 = 3;
//        index2 = 2;
//    }
//
//    rect[0] = temp_rect[index0];  // Left top coordinate
//    rect[1] = temp_rect[index1];  // Left bottom coordinate
//    rect[2] = temp_rect[index2];  // Right bottom coordinate
//    rect[3] = temp_rect[index3];  // Right top coordinate
//
//    if (rotated_rect.size.width < min_size ||
//        rotated_rect.size.height < min_size) {
//        return false;
//    } else {
//        return true;
//    }
//}

cv::RotatedRect expandBox(cv::Point2f temp[], float ratio) {
    ClipperLib::Path path = {
            {ClipperLib::cInt(temp[0].x), ClipperLib::cInt(temp[0].y)},
            {ClipperLib::cInt(temp[1].x), ClipperLib::cInt(temp[1].y)},
            {ClipperLib::cInt(temp[2].x), ClipperLib::cInt(temp[2].y)},
            {ClipperLib::cInt(temp[3].x), ClipperLib::cInt(temp[3].y)}};
    double area = ClipperLib::Area(path);
    double distance;
    double length = 0.0;
    for (int i = 0; i < 4; i++) {
        length = length + sqrtf(powf((temp[i].x - temp[(i + 1) % 4].x), 2) +
                                powf((temp[i].y - temp[(i + 1) % 4].y), 2));
    }

    distance = area * ratio / length;

    ClipperLib::ClipperOffset offset;
    offset.AddPath(path, ClipperLib::JoinType::jtRound,
                   ClipperLib::EndType::etClosedPolygon);
    ClipperLib::Paths paths;
    offset.Execute(paths, distance);

    std::vector <cv::Point> contour;
    for (int i = 0; i < paths[0].size(); i++) {
        contour.emplace_back(paths[0][i].X, paths[0][i].Y);
    }
    offset.Clear();
    return cv::minAreaRect(contour);
}

std::vector <std::vector<cv::Point2f>>
DBnet::Output(std::vector<float> &output, int height, int width, int src_h, int src_w) {
    cv::Mat map = cv::Mat::zeros(cv::Size(width, height), CV_8UC1);
    std::vector <std::vector<cv::Point2f>> boxes;
    for (int h = 0; h < height; ++h) {
        uchar *ptr = map.ptr(h);
        for (int w = 0; w < width; ++w) {
            ptr[w] = (output[h * width + w] > 0.3) ? 255 : 0;
        }
    }

    std::vector <std::vector<cv::Point>> contours;
    std::vector <cv::Vec4i> hierarcy;
    cv::findContours(map, contours, hierarcy, 1, 2);

    std::vector <cv::Rect> boundRect(contours.size());
    std::vector <cv::RotatedRect> box(contours.size());
    cv::Point2f rect[4];
    cv::Point2f order_rect[4];

    for (int i = 0; i < contours.size(); i++) {
        std::vector <cv::Point2f> points;
        cv::RotatedRect rotated_rect = cv::minAreaRect(cv::Mat(contours[i]));
        if (!get_mini_boxes(rotated_rect, rect, 5)) {
            std::cout << "box too small" << std::endl;
            continue;
        }

        // drop low score boxes
        float score = get_box_score(output, rect, inputW, inputH, 0.3);
        if (score < 0.5) {
            std::cout << "score too low =  " << score << ", threshold = " << 0.5 << std::endl;
            continue;
        }

        // Scaling the predict boxes depend on EXPANDRATIO
        cv::RotatedRect expandbox = expandBox(rect, 1.5);
        expandbox.points(rect);
        if (!get_mini_boxes(expandbox, rect, 5 + 2)) {
            continue;
        }

        // Restore the coordinates to the original image
        for (int k = 0; k < 4; k++) {
            order_rect[k] = rect[k];
            order_rect[k].x = int(order_rect[k].x / width * src_w);
            order_rect[k].y = int(order_rect[k].y / height * src_h);
            points.push_back(order_rect[k]);
        }
        if (points.size() > 0) boxes.push_back(points);
    }

    return boxes;
}

std::vector <std::vector<cv::Point2f>> DBnet::extract(const cv::Mat &img_input) {
    cudaSetDevice(device_id);
    inputW = img_input.rows;
    inputH = img_input.cols;
    float ratio = 1.0f;
    int max_len = inputW > inputH ? inputW : inputH;
    if (max_len > limit_side_len) {
        if (inputH > inputW) {
            ratio = float(limit_side_len) / float(inputH);
        } else {
            ratio = float(limit_side_len) / float(inputW);
        }
    }

    int resize_h = int(inputH * ratio) / 32 * 32;
    int resize_w = int(inputW * ratio) / 32 * 32;

    cv::Mat input;
    cv::resize(img_input, input, cv::Size(resize_w, resize_h));

    int single_image_size = channels * resize_h * resize_w;
    float *data = new float[single_image_size];
    std::vector<float> feats(1 * resize_h * resize_w);

    Input(input, data);

    context->setOptimizationProfile(0);

    nvinfer1::Dims4 inputDims{1, 3, resize_h, resize_w};
    nvinfer1::Dims out_dim = context->getBindingDimensions(0);
    context->setBindingDimensions(0, inputDims);
    int inputBufferSize = getMemorySize(inputDims);
    // int inputBufferSize = 1 * channels * resize_h * resize_w;
    CHECK(cudaMalloc(&bindings[0], inputBufferSize * sizeof(float)));

    auto outputDims = context->getBindingDimensions(1);
    int outputBufferSize = getMemorySize(outputDims);
    // int outputBufferSize = 1 * 1 * resize_h * resize_w;
    CHECK(cudaMalloc(&bindings[1], outputBufferSize * sizeof(float)));

    CHECK(cudaMemcpyAsync((float *) bindings[0], \
                data, single_image_size * sizeof(float), \
                cudaMemcpyHostToDevice, stream_t));
    bool status = context->enqueueV2(bindings, stream_t, nullptr);
    if (!status) {
        std::cout << "Enqueue failed" << std::endl;
    }

    cudaStreamSynchronize(stream_t);
    CHECK(cudaMemcpyAsync((float *) feats.data(), \
                (float *) bindings[1], \
                outputBufferSize * sizeof(float), cudaMemcpyDeviceToHost, stream_t));

    for (void *ptr : bindings) {
        cudaFree(ptr);
    }

    std::vector <std::vector<cv::Point2f>> boxes = Output(feats, resize_h, resize_w, img_input.rows, img_input.cols);

    delete[]data;
    return boxes;
}

DBnet::~DBnet() {
//    for (void *ptr : bindings) {
//        cudaFree(ptr);
//    }

    cudaStreamDestroy(stream_t);
}

REGISTER_TEXT_DETECTION_CLASS(DBnet);
